package com.xzy.mybatis.cipher.support.strategy.parse;

import cn.hutool.core.collection.CollectionUtil;
import com.alibaba.druid.sql.ast.SQLExpr;
import com.alibaba.druid.sql.ast.SQLStatement;
import com.alibaba.druid.sql.ast.statement.SQLUpdateSetItem;
import com.alibaba.druid.sql.ast.statement.SQLUpdateStatement;
import com.xzy.mybatis.cipher.support.registry.CryptoMetadataRegistry;
import com.xzy.mybatis.cipher.support.registry.DataCryptoStrategyRegistry;
import com.xzy.mybatis.cipher.support.strategy.crypto.CryptoStrategy;
import lombok.extern.slf4j.Slf4j;

import java.util.ArrayList;
import java.util.List;

/**
 * @author xzy.xiao
 * @since 2025/1/14  16:31
 */
@Slf4j
public class UpdateParseStrategy extends ParseStrategy {

    public UpdateParseStrategy(CryptoMetadataRegistry cryptoMetadataRegistry, DataCryptoStrategyRegistry dataCryptoStrategyRegistry) {
        super(cryptoMetadataRegistry, dataCryptoStrategyRegistry);
    }

    /**
     * UPDATE xxx SET xxx.aaa = ??? WHERE xxx.bbb = ???
     */
    @Override
    public List<Object> parse(String datasourceName, SQLStatement sqlStatement, List<Object> originalPropertyValueList) {
        SQLUpdateStatement sqlUpdateStatement = (SQLUpdateStatement) sqlStatement;
        String tableName = sqlUpdateStatement.getTableName().getSimpleName();
        String tableAliasName = sqlUpdateStatement.getTableSource().getAlias();

        if (!cryptoMetadataRegistry.isCryptoTable(datasourceName, tableName)) {
            return originalPropertyValueList;
        }

        List<Object> encryptedPropertyValueList = new ArrayList<>(originalPropertyValueList.size());

        // SET

        List<SQLUpdateSetItem> updateItemList = sqlUpdateStatement.getItems();
        if (updateItemList.isEmpty()) {
            return originalPropertyValueList;
        }

        int updateItemSize = updateItemList.size();
        for (int i = 0; i < updateItemSize; i++) {
            Object originalPropertyValue = originalPropertyValueList.get(i);
            SQLExpr updateColumn = updateItemList.get(i).getColumn();
            String updateColumnName = updateColumn.toString().replace(tableAliasName + ".", "");

            if (!cryptoMetadataRegistry.isCryptoColumn(datasourceName, tableName, updateColumnName)) {
                encryptedPropertyValueList.add(originalPropertyValue);
                continue;
            }

            CryptoStrategy cryptoStrategy = dataCryptoStrategyRegistry.getStrategy(datasourceName, tableName, updateColumnName);
            if (null == cryptoStrategy) {
                encryptedPropertyValueList.add(originalPropertyValue);
                continue;
            }

            String encryptedPropertyValue = cryptoStrategy.encrypt((String) originalPropertyValue);
            encryptedPropertyValueList.add(encryptedPropertyValue);
        }

        // WHERE
        List<String> whereColumnList = extractedWhere(sqlUpdateStatement.getWhere(), null);
        if (CollectionUtil.isEmpty(whereColumnList)) {
            return encryptedPropertyValueList;
        }

        for (int i = 0; i < whereColumnList.size(); i++) {

            Object originalPropertyValue = originalPropertyValueList.get(updateItemSize + i);
            String whereColumnName = whereColumnList.get(i).replace(tableAliasName + ".", "");

            if (!cryptoMetadataRegistry.isCryptoColumn(datasourceName, tableName, whereColumnName)) {
                encryptedPropertyValueList.add(originalPropertyValue);
                continue;
            }

            CryptoStrategy cryptoStrategy = dataCryptoStrategyRegistry.getStrategy(datasourceName, tableName, whereColumnName);

            if (null == cryptoStrategy) {
                encryptedPropertyValueList.add(originalPropertyValue);
                continue;
            }

            String encryptedPropertyValue = cryptoStrategy.encrypt((String) originalPropertyValue);
            encryptedPropertyValueList.add(encryptedPropertyValue);
        }

        log.debug("Parsed update sql.\n\t==> original property value:{}\n\t==> encrypted property value:{}", originalPropertyValueList, encryptedPropertyValueList);
        return encryptedPropertyValueList;
    }

}
